# coding: utf-8

# In[1]:


import time

import numpy as np
import tensorflow as tf

import tensorlayer as tl
from tensorlayer.layers import *

sess = tf.InteractiveSession()
X_train, y_train, X_test, y_test = tl.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False)


# In[2]:


def model(x, y_, reuse, is_train, batch_norm=False):
    
    #参数初始化函数
    W_init = tf.truncated_normal_initializer(stddev=5e-2)
    W_init2 = tf.truncated_normal_initializer(stddev=0.04)
    b_init2 = tf.constant_initializer(value=0.1)
    
    
    #定义模型及损失函数和准确度
    with tf.variable_scope("model", reuse=reuse):
        tl.layers.set_name_reuse(reuse)
        net = InputLayer(x, name='input')
        
        
        # 定义卷积1
        # Conv:
        #  act=tf.nn.relu
        #  shape=[5, 5, 3, 64],
        #  strides=[1, 1, 1, 1]
        # MaxPool:
        #  ksize=[1, 3, 3, 1]
        #  strides=[1, 2, 2, 1]
        
        if batch_norm:
            # 使用批规范化
            net = Conv2d(net, 64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='cnn1')
            net = BatchNormLayer(net, is_train, act=tf.nn.relu, name="batch1")
            net = MaxPool2d(net, (3, 3), (2, 2), padding='SAME', name='pool1')

        else:
            # 使用局部归一化
            net = Conv2d(net, 64, (5, 5), (1, 1), act=tf.nn.relu, padding='SAME', W_init=W_init, name='cnn1')
            net = MaxPool2d(net, (3, 3), (2, 2), padding='SAME', name='pool1')
            net = LocalResponseNormLayer(net, depth_radius=4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='norm1')

        
        # 定义卷积2
        # Conv:
        #  shape=[5, 5, 64, 64],
        #  strides=[1, 1, 1, 1], 
        # MaxPool:
        #  ksize=[1, 3, 3, 1]
        #  strides=[1, 2, 2, 1]
        
        if batch_norm:
            net = Conv2d(net, 64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='cnn2')
            net = BatchNormLayer(net, is_train, act=tf.nn.relu, name="batch2")
        else:
            net = Conv2d(net, 64, (5, 5), (1, 1), act=tf.nn.relu, padding='SAME', W_init=W_init, name='cnn2')
            net = LocalResponseNormLayer(net, depth_radius=4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='norm2')
        net = MaxPool2d(net, (3, 3), (2, 2), padding='SAME', name='pool2')

        
        # MLP分类模块
        net = FlattenLayer(net, name='flatten')  # output: (batch_size, 2304)
        net = DenseLayer(net, n_units=384, act=tf.nn.relu, W_init=W_init2, b_init=b_init2, name='d1relu')  # output: (batch_size, 384)
        net = DenseLayer(net, n_units=192, act=tf.nn.relu, W_init=W_init2, b_init=b_init2, name='d2relu')  # output: (batch_size, 192)
        net = DenseLayer(net, n_units=10, act=tf.identity, W_init=tf.truncated_normal_initializer(stddev=1 / 192.0), name='output')  # output: (batch_size, 10)
        y = net.outputs

        
        # 定义交叉熵损失
        ce = tl.cost.cross_entropy(y, y_, name='cost')
        
        # 定义L2惩罚
        L2 = 0
        for p in tl.layers.get_variables_with_name('relu/W', True, True):
            L2 += tf.contrib.layers.l2_regularizer(0.004)(p)
        cost = ce + L2

        # 准确度
        correct_prediction = tf.equal(tf.argmax(y, 1), y_)
        acc = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

        return net, cost, acc


# In[3]:


#数据增强
def distort_fn(x, is_train=False):
    x = tl.prepro.crop(x, 24, 24, is_random=is_train)
    if is_train:
        x = tl.prepro.flip_axis(x, axis=1, is_random=True)
        x = tl.prepro.brightness(x, gamma=0.1, gain=1, is_random=True)
    x = (x - np.mean(x)) / max(np.std(x), 1e-5)  # avoid values divided by 0
    return x


# In[4]:


def train():
    x = tf.placeholder(tf.float32, shape=[None, 24, 24, 3], name='x')
    y_ = tf.placeholder(
        tf.int64, shape=[
            None,
        ], name='y_')

    network, cost, _ = model(x, y_, False, is_train=True)
    _, cost_test, acc = model(x, y_, True, is_train=False)

    # train
    n_epoch = 50000
    learning_rate = 0.0001
    print_freq = 1
    batch_size = 2048

    train_params = network.all_params
    train_op = tf.train.AdamOptimizer(learning_rate, beta1=0.9, beta2=0.999, epsilon=1e-08, use_locking=False).minimize(cost, var_list=train_params)
    
    tl.layers.initialize_global_variables(sess)

    network.print_params(False)
    network.print_layers()

    print('   learning_rate: %f' % learning_rate)
    print('   batch_size: %d' % batch_size)

    for epoch in range(n_epoch):
        start_time = time.time()
        for X_train_a, y_train_a in tl.iterate.minibatches(X_train, y_train, batch_size, shuffle=True):
            X_train_a = tl.prepro.threading_data(X_train_a, fn=distort_fn, is_train=True)  # data augmentation for training
            sess.run(train_op, feed_dict={x: X_train_a, y_: y_train_a})

        if epoch + 1 == 1 or (epoch + 1) % print_freq == 0:
            print("Epoch %d of %d took %fs" % (epoch + 1, n_epoch, time.time() - start_time))
            test_loss, test_acc, n_batch = 0, 0, 0
            for X_test_a, y_test_a in tl.iterate.minibatches(X_test, y_test, batch_size, shuffle=False):
                X_test_a = tl.prepro.threading_data(X_test_a, fn=distort_fn, is_train=False)  # central crop
                err, ac = sess.run([cost_test, acc], feed_dict={x: X_test_a, y_: y_test_a})
                test_loss += err
                test_acc += ac
                n_batch += 1
            print("   test loss: %f" % (test_loss / n_batch))
            print("   test acc: %f" % (test_acc / n_batch))


# In[5]:
import sys
if __name__ == '__main__':
    if (len(sys.argv) >= 2 and sys.argv[1] == 'f'):
        print("force_cpu")
        tf.device("/cpu:0")
    train()


